use crate::Rule;
use pest::iterators::Pair;

pub struct Formatter {
    s: String,
    indent_level: usize,
    indent_size: usize,
    comp_unit_count: usize
}

impl Formatter {
    pub fn new() -> Self{
        Self {
            s: String::new(),
            indent_level: 0,
            indent_size: 4,
            comp_unit_count: 0
        }
    }

    pub fn print(&self) {
        print!("{}", self.s)
    }

    fn append(&mut self, s: &str) -> &mut Self {
        self.s.push_str(s);
        self
    }
    fn indent(&mut self) -> &mut Self {
        self.s.push_str(" ".repeat(self.indent_level * self.indent_size).as_mut_str());
        self
    }
    fn add_indent(&mut self) -> &mut Self {
        self.indent_level += 1;
        self
    }
    fn del_indent(&mut self) -> &mut Self {
        if self.indent_level > 0 {
            self.indent_level -= 1;
        }
        self
    }

    pub fn format(&mut self, pair: Pair<Rule>){
        for inner_pair in pair.into_inner() {
            match inner_pair.as_rule() {
                Rule::CompUnit => { self.format_comp_unit(inner_pair); },
                _ => ()
            };
        }
    }


    fn format_comp_unit(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::Decl => { self.format_decl(p); }
                Rule::FuncDef => {
                    self.format_func_def(p);
                }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_decl(&mut self, pair: Pair<Rule>) -> &mut Self {
        self.comp_unit_count += 1;
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::ConstDecl => self.format_const_decl(p),
                Rule::VarDecl => self.format_var_decl(p),
                other => panic!("unexpected rule {other:?}")
            };
        }
        self.append("\n")
    }

    fn format_const_decl(&mut self, pair: Pair<Rule>) -> &mut Self {
        self.append("const ");
        let mut i = 0;
        for p in pair.into_inner() {
            if i == 0 { self.format_btype(p).append(" "); }
            else if i == 1 { self.format_const_def(p); }
            else {
                self.append(", ").format_const_def(p);
            }

            i += 1;
        }
        self.append(";")
    }

    fn format_btype(&mut self, pair: Pair<Rule>) -> &mut Self { self.append(pair.as_str().trim()) }

    fn format_const_def(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::Ident => { self.format_ident(p); }
                Rule::ConstExp => { self.append("[").format_const_exp(p).append("]"); }
                Rule::ConstInitVal => { self.append(" = ").format_const_init_val(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_const_init_val(&mut self, pair: Pair<Rule>) -> &mut Self {
        if let None = pair.clone().into_inner().next() {
            return self.append("{}");
        }
        let mut i = 0;
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::ConstExp => { self.format_const_exp(p); }
                Rule::ConstInitVal => {
                    if i == 0 { self.append("{").format_const_init_val(p); }
                    else { self.append(", ").format_const_init_val(p); }
                    i += 1;
                }
                other => panic!("unexpected rule {other:?}")
            }
        }
        if i != 0 { self.append("}") }
        else { self }
    }

    fn format_var_decl(&mut self, pair: Pair<Rule>) -> &mut Self {
        let mut i = 0;
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::BType => { self.format_btype(p).append(" "); }
                Rule::VarDef => {
                    if i == 0 { self.format_var_def(p); }
                    else { self.append(", ").format_var_def(p); }
                    i += 1;
                }
                other => panic!("unexpected rule {other:?}")
            };
        }
        self.append(";")
    }

    fn format_var_def(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::Ident => { self.format_ident(p); }
                Rule::ConstExp => { self.append("[").format_const_exp(p).append("]"); }
                Rule::InitVal => { self.append(" = ").format_init_val(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_init_val(&mut self, pair: Pair<Rule>) -> &mut Self {
        if let None = pair.clone().into_inner().next() {
            return self.append("{}");
        }
        let mut i = 0;
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::Exp => { self.format_exp(p); }
                Rule::InitVal => {
                    if i == 0 { self.append("{").format_init_val(p); }
                    else { self.append(", ").format_init_val(p); }
                    i += 1;
                }
                other => panic!("unexpected rule {other:?}")
            }
        }
        if i != 0 { self.append("}") }
        else { self }
    }

    fn format_func_def(&mut self, pair: Pair<Rule>) -> &mut Self {
        self.comp_unit_count += 1;
        if self.comp_unit_count > 1 { self.append("\n"); }
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::FuncType => { self.format_func_type(p).append(" "); }
                Rule::Ident => { self.format_ident(p).append("("); }
                Rule::FuncFParams => { self.format_func_fparams(p); }
                Rule::Block => { self.append(") ").format_block(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self.append("\n")
    }

    fn format_func_type(&mut self, pair: Pair<Rule>) -> &mut Self { self.append(pair.as_str().trim()) }

    fn format_func_fparams(&mut self, pair: Pair<Rule>) -> &mut Self {
        let mut i = 0;
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::FuncFParam => {
                    if i == 0 { self.format_func_fparam(p); }
                    else { self.append(", ").format_func_fparam(p); }
                    i += 1;
                }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_func_fparam(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::BType => { self.format_btype(p).append(" "); }
                Rule::Ident => { self.format_ident(p); }
                Rule::EmptyBrackets => { self.format_empty_brackets(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_empty_brackets(&mut self, _pair: Pair<Rule>) -> &mut Self {
        self.append("[]")
    }

    fn format_block(&mut self, pair: Pair<Rule>) -> &mut Self {
        self.append("{\n").add_indent();
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::BlockItem => { self.indent().format_block_item(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self.del_indent().indent().append("}")
    }

    fn format_block_item(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::Decl => { self.format_decl(p); }
                Rule::Stmt => { self.format_stmt(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_stmt(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::AssignStmt => { self.format_assign_stmt(p).append("\n"); }
                Rule::ExpStmt => { self.format_exp_stmt(p).append("\n"); }
                Rule::Block => { self.format_block(p).append("\n"); }
                Rule::IfStmt => { self.format_if_stmt(p); }
                Rule::WhileStmt => { self.format_while_stmt(p); }
                Rule::BreakStmt => { self.format_break_stmt(p).append("\n"); }
                Rule::ContinueStmt => { self.format_continue_stmt(p).append("\n"); }
                Rule::ReturnStmt => { self.format_return_stmt(p).append("\n"); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_assign_stmt(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::LVal => { self.format_lval(p).append(" = "); }
                Rule::Exp => { self.format_exp(p).append(";"); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_exp_stmt(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::Exp => { self.format_exp(p); } 
                other => panic!("unexpected rule {other:?}")
            }
        }
        self.append(";")
    }

    fn format_if_stmt(&mut self, pair: Pair<Rule>) -> &mut Self {
        self.append("if ");
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::Cond => { self.append("(").format_cond(p).append(")"); }
                Rule::Stmt => { 
                    if let Rule::Block = p.clone().into_inner().next().unwrap().as_rule() {
                        self.append(" ").format_stmt(p);
                    }
                    else { 
                        self.add_indent().append("\n").indent(); 
                        self.format_stmt(p).del_indent();
                    }

                }
                Rule::Else => { self.format_else(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_else(&mut self, pair: Pair<Rule>) -> &mut Self {
        self.indent().append("else");
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::Stmt => {
                    if let Rule::Block | Rule::IfStmt = p.clone().into_inner().next().unwrap().as_rule() {
                        self.append(" ").format_stmt(p);
                    }
                    else { 
                        self.add_indent().append("\n").indent(); 
                        self.format_stmt(p).del_indent();
                    }
                }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }


    fn format_while_stmt(&mut self, pair: Pair<Rule>) -> &mut Self {
        self.append("while ");
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::Cond => { self.append("(").format_cond(p).append(")"); }
                Rule::Stmt => { 
                    if let Rule::Block = p.clone().into_inner().next().unwrap().as_rule() {
                        self.append(" ").format_stmt(p);
                    }
                    else { 
                        self.add_indent().append("\n").indent(); 
                        self.format_stmt(p).del_indent();
                    }

                }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_break_stmt(&mut self, _pair: Pair<Rule>) -> &mut Self { self.append("break;") }

    fn format_continue_stmt(&mut self, _pair: Pair<Rule>) -> &mut Self { self.append("continue;") }

    fn format_return_stmt(&mut self, pair: Pair<Rule>) -> &mut Self {
        match pair.into_inner().next() {
            Some(p) => {
                if let Rule::Exp = p.as_rule() {
                    self.append("return ").format_exp(p).append(";");
                }
                else { panic!("unexpected rule {:?}", p.as_rule()) }

            }
            None => { return self.append("return;"); }
        };
        self
    }

    fn format_exp(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::AddExp => { self.format_add_exp(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_cond(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::LOrExp => { self.format_lor_exp(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_lval(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::Ident => { self.format_ident(p); }
                Rule::Exp => { self.append("[").format_exp(p).append("]"); } 
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_primary_exp(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::Exp => { self.append("(").format_exp(p).append(")"); } 
                Rule::LVal => { self.format_lval(p); }
                Rule::Number => { self.format_number(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_number(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::IntConst => { self.format_int_const(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_unary_exp(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::Ident => { self.format_ident(p); }
                Rule::L_PAREN => { self.append("("); }
                Rule::R_PAREN => { self.append(")"); }
                Rule::FuncRParams => { self.format_func_rparams(p); }
                Rule::PrimaryExp => { self.format_primary_exp(p); }
                Rule::UnaryOp => { self.format_unary_op(p); }
                Rule::UnaryExp => { self.format_unary_exp(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_unary_op(&mut self, pair: Pair<Rule>) -> &mut Self { self.append(pair.as_str().trim()) }

    fn format_func_rparams(&mut self, pair: Pair<Rule>) -> &mut Self {
        let mut i = 0;
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::Exp => {
                    if i == 0 { self.format_exp(p); }
                    else { self.append(", ").format_exp(p); }
                    i += 1;
                }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_mul_exp(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::UnaryExp => { self.format_unary_exp(p); }
                Rule::MulOp => { self.format_mul_op(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_add_exp(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::MulExp => { self.format_mul_exp(p); }
                Rule::AddOp => { self.format_add_op(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_rel_exp(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::AddExp => { self.format_add_exp(p); }
                Rule::RelOp => { self.format_rel_op(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_eq_exp(&mut self, pair: Pair<Rule>) -> &mut Self {
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::RelExp => { self.format_rel_exp(p); }
                Rule::EqOp => { self.format_eq_op(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_land_exp(&mut self, pair: Pair<Rule>) -> &mut Self {
        let mut i = 0;
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::EqExp => {
                    if i == 0 { self.format_eq_exp(p); }
                    else { self.append(" && ").format_eq_exp(p); }
                    i += 1;
                }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_lor_exp(&mut self, pair: Pair<Rule>) -> &mut Self {
        let mut i = 0;
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::LAndExp => {
                    if i == 0 { self.format_land_exp(p); }
                    else { self.append(" || ").format_land_exp(p); }
                    i += 1;
                }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_const_exp(&mut self, pair: Pair<Rule>) -> &mut Self{
        for p in pair.into_inner() {
            match p.as_rule() {
                Rule::AddExp => { self.format_add_exp(p); }
                other => panic!("unexpected rule {other:?}")
            }
        }
        self
    }

    fn format_ident(&mut self, pair: Pair<Rule>) -> &mut Self { self.append(pair.as_str().trim()) }

    fn format_int_const(&mut self, pair: Pair<Rule>) -> &mut Self { self.append(pair.as_str().trim() ) }

    fn format_mul_op(&mut self, pair: Pair<Rule>) -> &mut Self { self.append(&format!(" {} ", pair.as_str().trim())) }

    fn format_add_op(&mut self, pair: Pair<Rule>) -> &mut Self { self.append(&format!(" {} ", pair.as_str().trim())) }
    fn format_rel_op(&mut self, pair: Pair<Rule>) -> &mut Self { self.append(&format!(" {} ", pair.as_str().trim())) }
    fn format_eq_op(&mut self, pair: Pair<Rule>) -> &mut Self { self.append(&format!(" {} ", pair.as_str().trim())) }

}
